# aegaeon_proto/server.py
import os
import time
import json
import uuid
import asyncio
import logging
from typing import Dict, Set
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
import redis.asyncio as aioredis

# config
REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
REQ_QUEUE = "aegaeon:requests"             # fallback global queue
NODE_QUEUE_PREFIX = "aegaeon:req:node:"   # per-node queue: + node_id
PROCESSING_LIST = "aegaeon:processing"
RESULT_HASH = "aegaeon:results"
MODEL_LOC_PREFIX = "aegaeon:model_locations:"  # model -> set(node_id)
RESULT_PUB_CHANNEL = "aegaeon:results:channel"
NODE_REGISTRY = "aegaeon:nodes"           # node heartbeat/info

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("aegaeon.server")

app = FastAPI()
redis = aioredis.from_url(REDIS_URL)

# WebSocket subscriptions: map req_id -> set of websocket connections
WS_SUBSCRIBERS: Dict[str, Set[WebSocket]] = {}
WS_LOCK = asyncio.Lock()

def _node_queue(node_id: str) -> str:
    return NODE_QUEUE_PREFIX + node_id

async def _pick_node_for_model(model: str):
    """Return a node id that likely has model prefetched (round-robin among set),
       or None if not found."""
    key = MODEL_LOC_PREFIX + model
    nodes = await redis.smembers(key)
    if not nodes:
        return None
    # pick first (could be improved to round robin)
    # convert set to sorted list for stability
    nodes_list = sorted(nodes)
    return nodes_list[0]

async def _push_to_node_queue(node_id: str, meta: dict):
    q = _node_queue(node_id)
    await redis.rpush(q, json.dumps(meta))

async def _push_global(meta: dict):
    await redis.rpush(REQ_QUEUE, json.dumps(meta))

@app.post("/v1/generate")
async def generate(req: Request):
    body = await req.json()
    model = body.get("model")
    prompt = body.get("prompt")
    callback = body.get("callback_url")  # optional HTTP callback (worker will POST result)
    client_subscribe = body.get("subscribe_ws", False)  # if true, client will also open WS and subscribe to req_id
    if not model or not prompt:
        return {"error": "model and prompt required"}, 400
    req_id = body.get("req_id") or f"r-{int(time.time()*1000)}-{uuid.uuid4().hex[:6]}"
    created = time.time()
    meta = {
        "req_id": req_id,
        "model": model,
        "prompt": prompt,
        "created": created,
        "callback_url": callback,
        "client_subscribe": client_subscribe,
        "client_meta": body.get("client_meta", {}),
    }

    # Try to route to node that already has model cached
    node = await _pick_node_for_model(model)
    if node:
        await _push_to_node_queue(node, meta)
        log.info(f"Queued request {req_id} for model {model} to node {node} (model-local)")
    else:
        # push to global queue for workers to consume (they may prefetch)
        await _push_global(meta)
        log.info(f"Queued request {req_id} for model {model} to global queue")

    return {"req_id": req_id, "status": "accepted", "queued_at": created}

@app.get("/v1/result")
async def get_result(req_id: str):
    if not req_id:
        return {"error": "req_id required"}, 400
    res = await redis.hget(RESULT_HASH, req_id)
    if res is None:
        return {"req_id": req_id, "status": "pending"}
    try:
        obj = json.loads(res)
    except Exception:
        obj = {"raw": res}
    return {"req_id": req_id, "status": "done", "result": obj}

@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    """
    WebSocket endpoint. Client should connect and then send JSON:
      {"action":"subscribe","req_id":"r-..."}
    Server forwards result messages published on Redis channel to subscribed clients.
    """
    await ws.accept()
    sub_req_ids = set()
    try:
        while True:
            txt = await ws.receive_text()
            try:
                msg = json.loads(txt)
            except Exception:
                await ws.send_text(json.dumps({"error": "invalid json"}))
                continue
            action = msg.get("action")
            if action == "subscribe":
                rid = msg.get("req_id")
                if not rid:
                    await ws.send_text(json.dumps({"error": "req_id required for subscribe"})); continue
                async with WS_LOCK:
                    WS_SUBSCRIBERS.setdefault(rid, set()).add(ws)
                    sub_req_ids.add(rid)
                await ws.send_text(json.dumps({"subscribed": rid}))
            elif action == "unsubscribe":
                rid = msg.get("req_id")
                async with WS_LOCK:
                    if rid in WS_SUBSCRIBERS and ws in WS_SUBSCRIBERS[rid]:
                        WS_SUBSCRIBERS[rid].remove(ws)
                if rid in sub_req_ids: sub_req_ids.remove(rid)
                await ws.send_text(json.dumps({"unsubscribed": rid}))
            else:
                await ws.send_text(json.dumps({"error": "unknown action"}))
    except WebSocketDisconnect:
        # cleanup
        async with WS_LOCK:
            for rid in sub_req_ids:
                if rid in WS_SUBSCRIBERS and ws in WS_SUBSCRIBERS[rid]:
                    WS_SUBSCRIBERS[rid].remove(ws)
        return

# background task: subscribe to Redis pubsub and forward to connected WebSocket clients
@app.on_event("startup")
async def startup_event():
    asyncio.create_task(_redis_result_forwarder())

async def _redis_result_forwarder():
    sub = redis.pubsub()
    await sub.subscribe(RESULT_PUB_CHANNEL)
    log.info("Subscribed to result channel for WS forwarding")
    async for message in sub.listen():
        # message: dict with type, channel, data
        if message is None:
            await asyncio.sleep(0.01); continue
        typ = message.get("type")
        if typ != "message": continue
        data = message.get("data")
        if not data:
            continue
        try:
            obj = json.loads(data)
            req_id = obj.get("req_id")
            if not req_id:
                continue
        except Exception:
            continue
        # forward to WS subscribers (if any)
        async with WS_LOCK:
            conns = list(WS_SUBSCRIBERS.get(req_id, set()))
        if conns:
            payload = json.dumps({"req_id": req_id, "result": obj})
            for ws in conns:
                try:
                    await ws.send_text(payload)
                except Exception:
                    # ignore - client may have disconnected; cleanup happens on disconnect
                    pass

